from ui.components.base.widget import TgBaseWidget
from ui.components.layouts.vbox_layout import TgVBoxLayout
from ui.components.layouts.hbox_layout import TgHBoxLayout
from ui.components.inputs.label import TgLabel
from ui.components.buttons.button import TgButton
from ui.components.composite.plot_widget import TgPlotWidget
from ui.components.charts.html_table import TgHtmlTable
from ui.components.charts.bar_chart import TgBarChart
from ui.components.inputs.date_edit import TgDateEdit
import pandas as pd
from datetime import datetime, timedelta
from PySide6.QtWidgets import QScrollArea
from loguru import logger
from util import get_latest_trading_date


class IndexDataView(TgBaseWidget):
    
    def __init__(self):
        super().__init__()
        self.layout = TgVBoxLayout(self)
        self._setup_ui()
        
    def _setup_ui(self):
        scroll_area = QScrollArea()
        scroll_area.setWidgetResizable(True)

        content_widget = TgBaseWidget()
        content_layout = TgVBoxLayout(content_widget)

        self._create_table(content_layout)
        content_layout.addSpacing(10)
        
        self._create_date_controls(content_layout)
        content_layout.addSpacing(10)
        
        self._create_charts(content_layout)
        content_layout.addStretch(1)

        scroll_area.setWidget(content_widget)
        self.layout.addWidget(scroll_area)

    def _create_table(self, layout):
        headers = ["指数代码", "指数名称", "日内走势", "近一日", "近一周", "近三周", "近一月", "近2月", "近3月", "近1年", "年初至今"]
        self.table = TgHtmlTable(headers=headers, cell_size='small')
        self.table.setFixedHeight(1200)
        
        self.table.set_profit_loss_style(list(range(3, 11)), use_gradient=True, color_type='background')
        self.table.set_chart_column(2, {
            'width': 80, 
            'height': 30, 
            'color_style': 'stock_profit'
        })
        self.table.enable_sorting(enable=True)
        
        layout.addWidget(self.table)

    def _create_date_controls(self, layout):
        controls_layout = TgHBoxLayout()
        controls_layout.addWidget(TgLabel("起始时间:"))
        self.start_date_input = TgDateEdit()
        controls_layout.addWidget(self.start_date_input)

        controls_layout.addWidget(TgLabel("结束时间:"))
        self.end_date_input = TgDateEdit()
        controls_layout.addWidget(self.end_date_input)

        self.query_button = TgButton("查询")
        controls_layout.addWidget(self.query_button)
        controls_layout.addStretch(1)

        layout.addLayout(controls_layout)

    def _create_charts(self, layout):
        self.chart_volume_ratio = TgBarChart(self)
        self.chart_volume_ratio.setFixedHeight(300)
        self.chart_volume_ratio.setObjectName("volume_ratio_chart")
        layout.addWidget(self.chart_volume_ratio)
        layout.addSpacing(10)

        self.charts = {}
        chart_configs = [
            ("sz50_vs_zz2000", "大小比价(上证50/中证2000)"),
            ("csi300_vs_csi500", "大小比价(沪深300/中证500)"),
            ("csi300_vs_csi1000", "大小比价(沪深300/中证1000)"),
            ("csi300_vs_zz2000", "大小比价(沪深300/中证2000)"),
            ("cumulative_return", "主要指数累计收益率 (基准: 100)"),
            ("style_return", "风格指数累计收益率 (基准: 100)"),
            ("asset_class_return", "大类资产累计收益率 (基准: 100)"),
            ("sector_return", "行业板块累计收益率 (基准: 100)")
        ]
        
        for chart_id, title in chart_configs:
            chart = TgPlotWidget(title=title, time_axis=True)
            if "cumulative_return" in chart_id or "return" in chart_id:
                chart.setFixedHeight(400)
            else:
                chart.setFixedHeight(300)
            chart.setObjectName(chart_id)
            chart.set_mouse_interaction(enabled=True)
            self.charts[chart_id] = chart
            layout.addWidget(chart)

    def set_default_dates(self, start_date: str, end_date: str):
        # 获取最近一个交易日作为结束日期
        latest_trading_date = get_latest_trading_date()
        if latest_trading_date:
            end_dt = datetime.strptime(latest_trading_date, '%Y-%m-%d').date()
            logger.info(f"使用最近交易日作为结束日期: {latest_trading_date}")
        else:
            # 如果获取失败，使用传入的结束日期
            end_dt = datetime.strptime(end_date, '%Y-%m-%d').date()
            logger.warning(f"获取最近交易日失败，使用默认结束日期: {end_date}")
        
        start_dt = datetime.strptime(start_date, '%Y-%m-%d').date()
        
        self.end_date_input.set_python_date(end_dt)
        self.start_date_input.set_python_date(start_dt)

    def get_date_range(self):
        start_date = self.start_date_input.get_date_string("yyyy-MM-dd")
        end_date = self.end_date_input.get_date_string("yyyy-MM-dd")
        return start_date, end_date

    def update_table_data(self, data):
        self.table.update_data(data=data)

    def set_query_loading(self, is_loading: bool):
        self.query_button.setEnabled(not is_loading)
        self.query_button.setText("查询中..." if is_loading else "查询")

    def update_ratio_chart(self, chart_name: str, chart_data: dict):
        if chart_name in self.charts:
            chart = self.charts[chart_name]
            chart.clear()
            chart.add_line_chart(
                x=chart_data['x'],
                y=chart_data['y'],
                name=chart_data['name'],
                color='#9b59b6'
            )
            chart.auto_range()

    def update_cumulative_return_chart(self, chart_name: str, indices: dict, index_data: dict):
        chart_map = {
            "main": "cumulative_return",
            "style": "style_return", 
            "asset_class": "asset_class_return",
            "sector": "sector_return"
        }
        
        chart_id = chart_map.get(chart_name)
        if not chart_id or chart_id not in self.charts:
            return
            
        chart = self.charts[chart_id]
        chart.clear()
        
        colors = ['#1f77b4', '#ff7f0e', '#2ca02c', '#d62728', '#9467bd', '#8c564b', '#e377c2', '#7f7f7f', '#bcbd22', '#17becf']
        color_index = 0
        
        for code, name in indices.items():
            df = index_data.get(code)
            if df is not None and not df.empty:
                df = df.sort_values('date').reset_index(drop=True)
                if not df.empty:
                    df['cumulative_return'] = (df['close'] / df['close'].iloc[0]) * 100
                    chart.add_line_chart(
                        x=df['date'].tolist(),
                        y=df['cumulative_return'].tolist(),
                        name=name,
                        color=colors[color_index % len(colors)]
                    )
                    color_index += 1
                    logger.info(f"成功加载 {name} 累计收益率数据")
        chart.auto_range()

    def update_volume_ratio_chart(self, indices: dict, index_data: dict):
        valid_dates = None
        volume_data = {}
        
        for code, name in indices.items():
            df = index_data.get(code)
            if df is not None and not df.empty and 'money' in df.columns:
                df = df.sort_values('date').reset_index(drop=True)
                df = df[df['money'] > 0].dropna(subset=['money'])
                if not df.empty:
                    volume_data[code] = df
                    if valid_dates is None:
                        valid_dates = set(df['date'].dt.strftime('%Y-%m-%d'))
                    else:
                        valid_dates &= set(df['date'].dt.strftime('%Y-%m-%d'))
        
        if not volume_data or not valid_dates:
            logger.warning("没有足够的成交额数据来绘制占比图表")
            return
        
        sorted_dates = sorted(list(valid_dates))
        
        daily_ratios = {}
        for code, name in indices.items():
            if code in volume_data:
                daily_ratios[code] = []
                df = volume_data[code]
                for date_str in sorted_dates:
                    date_data = df[df['date'].dt.strftime('%Y-%m-%d') == date_str]
                    if not date_data.empty:
                        daily_ratios[code].append(float(date_data.iloc[0]['money']))
                    else:
                        daily_ratios[code].append(0.0)
        
        ratio_series_data = []
        for code, name in indices.items():
            if code in daily_ratios:
                ratio_data = []
                for i in range(len(sorted_dates)):
                    total_volume = sum(daily_ratios[c][i] for c in daily_ratios.keys())
                    if total_volume > 0:
                        ratio = daily_ratios[code][i] / total_volume
                        ratio_data.append(round(ratio, 4))
                    else:
                        ratio_data.append(0.0)
                
                ratio_series_data.append({
                    "name": name,
                    "data": ratio_data
                })
        
        df = pd.DataFrame()
        df['date'] = sorted_dates
        y_cols = []
        for series in ratio_series_data:
            y_cols.append(series['name'])
            df[series['name']] = series['data']

        self.chart_volume_ratio.set_title("指数成交额占比")
        self.chart_volume_ratio.set_data(df, x_axis_column='date', y_axis_columns=y_cols)
        self.chart_volume_ratio.set_stacked(True)
        self.chart_volume_ratio.set_labels(False)
        self.chart_volume_ratio.set_datazoom(False)
        self.chart_volume_ratio.set_compact_mode(True)
        logger.info("成功加载成交额占比堆叠柱状图")

    def clear_all_charts(self):
        for chart in self.charts.values():
            chart.clear()
        self.chart_volume_ratio.clear()